import requests
import json
import re 
import pymongo

#连接数据库，后续存储爬取到的数据
class MongoDBLink(object):
    def __init__(self, host='127.0.0.1', port=27017):
        # 连接mongodb

        self.host = host
        self.port = port
        self.client = pymongo.MongoClient(host=self.host, port=self.port)
        # 切换数据库
        self.db = self.client['CNNVD_vul']
        # self.db = self.client.get_database('db_name')
        # self.db = self.client.db_name

    # 往数据库中添加一条内容
    def add_one(self, data):
        #为了给不同名称的集合插入不同的数据，我们动态调用插入函数，
        result = self.db.vuldb.insert_one(data)
        # inserted_id返回插入记录的id
        # print(result.inserted_id)

#创建数据库对象
mongo_py = MongoDBLink()


#定义获取漏洞的全部信息的函数。ddict至少包含：ddict={"cnnvdCode": "","id": "","vulType": ""}
def get_all_vuldata(ddict):
    #单个漏洞详细信息页的请求头
    headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/95.0.4638.69 Safari/537.36',
        'Host': 'www.cnnvd.org.cn',
        'Content-Type': 'application/json',
        'Origin': 'https://www.cnnvd.org.cn',
        'Referer': 'https://www.cnnvd.org.cn/home/loophole',
    }
    #参数，也可以称为payload（载荷）
    params = ddict
    request_url="https://www.cnnvd.org.cn/web/cnnvdVul/getCnnnvdDetailOnDatasource"
    r = requests.post(request_url,data = json.dumps(params),headers = headers)
    #print(r.text)
   
    #正则表达式提取需要信息，切去冗余信息
    finddata=re.compile(r'"cnnvdDetail":(.*?),"receviceVulDetail"')
    data=re.findall(finddata,r.text)
    #此时data的样式为['{"id":null......."varchar1":"命令注入"}']
    #data[0]为'{"id":null......."varchar1":"命令注入"}'，是一个字符串，包含了单个漏洞的所有信息
    print(data[0],type(data[0]))

    #去除data[0]字符串中的空值字符串，使其可以正常转化为字典，方便数据库存储
    if (":null" in data[0]):
        data[0]=data[0].replace(":null",":\'null\'")
    #处理好null字符的字符串，转化为字典
    datavul=eval(data[0])
    #print(type(datavul),type(ddict))

    #声明这两个数据类型为字典，否则不可以进行updata操作
    vuldata1=dict(datavul)
    vuldata2=dict(ddict)
    #print(type(vuldata1),type(vuldata2))

    #将传入函数的，单个漏洞的部分信息ddict添加到单个漏洞的详细信息中作为补充
    vuldata1.update(vuldata2)
    #print(vuldata1,type(vuldata1))

    #存储到数据库中
    mongo_py.add_one(vuldata1)

#包含网页的请求头等信息，帮助伪造post请求
headers = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/90.0.4430.212 Safari/537.36',
        'Host': 'www.cnnvd.org.cn',
        'Content-Type': 'application/json',
        'Origin': 'https://www.cnnvd.org.cn',
        'Referer': 'https://www.cnnvd.org.cn/home/loophole',
}

#循环爬取，从首页到尾页
for j in range(1,22446):

    #参数，"pageIndex":j是循环爬取的关键
    params = {"pageIndex":j,"pageSize":10,"keyword":"","hazardLevel":"","vulType":"","vendor":"","product":"","dateType":""}
    request_url="https://www.cnnvd.org.cn/web/homePage/cnnvdVulList"
    r = requests.post(request_url,data = json.dumps(params),headers = headers)
    # http接口测试返回码statusCode含义:200（成功）  服务器已成功处理了请求。通常，这表示服务器提供了请求的网页。其余返回结果还有404等

    #正则表达式提取需要信息，切去冗余信息
    finddata=re.compile(r'"records":(.*?),"pageIndex"')
    data=re.findall(finddata,r.text)

    #此时data的格式为['[{"id":"f3c9e7b9b...........其他","vulType":"0"}]'],len=1,包含该页所有漏洞
    #此时data[0]的格式为字符串：'[{"id":"f3c9e7b9b......},.....,{.....其他","vulType":"0"}]'

    # print(data[0])

    #替换字典中出现的null，将空值替换为字符串“null”，以便后续正确转化为字典格式
    if (":null" in data[0]):
        data[0]=data[0].replace(":null",":\'null\'")
    #将data[0]转化为列表，列表中的元素是单个漏洞数据构成的字典
    d=eval(data[0])
    #循环爬取详细信息
    for i in d:
            get_all_vuldata(i)   
    print(j)

#hazardLevel: 1(超危),2(高危),3(中危),4(低危),0(无定级)
